local backends = require 'scripts.core.types.trueskill.backends'
local CDF, PDF, PPF = backends.cdf, backends.pdf, backends.ppf
local factorgraph = require 'scripts.core.types.trueskill.factorgraph'
local Variable = factorgraph.Variable
local PriorFactor = factorgraph.PriorFactor
local LikelihoodFactor = factorgraph.LikelihoodFactor
local SumFactor = factorgraph.SumFactor
local TruncateFactor = factorgraph.TruncateFactor
local mathematics = require 'scripts.core.types.trueskill.mathematics'
local Gaussian = mathematics.Gaussian
local Matrix = mathematics.Matrix

local MU = 25
local SIGMA = MU / 3
local BETA = SIGMA / 2
local TAU = SIGMA / 100
local DRAW_PROBABILITY = 0.10
local DELTA = 0.0001

local global_env_var 

-- 定义TrueSkill类
local TrueSkill = {}
TrueSkill.__index = TrueSkill

-- 构造函数
function TrueSkill:new(mu, sigma, beta, tau, draw_probability, backend)
    local self = setmetatable({}, TrueSkill)
    self.mu = mu or MU
    self.sigma = sigma or SIGMA
    self.beta = beta or BETA
    self.tau = tau or TAU
    self.draw_probability = draw_probability or DRAW_PROBABILITY
    self.backend = backend
    if type(backend) == "table" and #backend == 3 then
        self.cdf, self.pdf, self.ppf = table.unpack(backend)
    else
        self.cdf, self.pdf, self.ppf = CDF, PDF, PPF
    end
    return self
end

local function setup(options)
    options = options or {}
    local env = options.env or TrueSkill.new(options)
    global_env_var = env
    return env
end

local function global_env()
    if global_env_var == nil then
        setup()
    end
    return global_env_var
end

-- 定义Rating类
local Rating = {}
Rating.__index = Rating

-- 构造函数
function Rating:new(mu, sigma)
    if type(mu) == "table" and #mu == 2 then
        mu, sigma = table.unpack(mu)
    elseif type(mu) == "table" and mu.mu and mu.sigma then
        mu, sigma = mu.mu, mu.sigma
    end
    if mu == nil then
        mu = global_env().mu
    end
    if sigma == nil then
        sigma = global_env().sigma
    end
    local self = setmetatable(Gaussian(mu,sigma),Rating)
    self.mu = mu
    self.sigma = sigma
    return self
end

function Rating:__tostring()
    return string.format("Rating(mu=%.3f, sigma=%.3f)", self.mu, self.sigma)
end

setmetatable(Rating, {
    __index = Gaussian,
    __call = Rating.new,
})

-- 定义create_rating方法
function TrueSkill:create_rating(mu, sigma)
    if mu == nil then
        mu = self.mu
    end
    if sigma == nil then
        sigma = self.sigma
    end
    return Rating(mu, sigma)
end

function TrueSkill:validate_rating_groups(rating_groups)
    if #rating_groups < 2 then
        error('Need multiple rating groups')
    end
    for _, group in ipairs(rating_groups) do
        if #group == 0 then
            error('Each group must contain multiple ratings')
        end
    end
    local group_types = {}
    for _, group in ipairs(rating_groups) do
            group_types[getmetatable(group) or type(group)] = true
    end
    if next(group_types, next(group_types)) then
        error('All groups should be same type')
    elseif group_types[Rating] then
        error('Rating cannot be a rating group')
    end
    return rating_groups
end

function TrueSkill:validate_weights(weights, rating_groups)
    if weights == nil then
        weights = {}
        for _, g in ipairs(rating_groups) do
            table.insert(weights, {})
            for _ = 1, #g do
                table.insert(weights[#weights], 1)
            end
        end
    end 
    return weights 
end 

local function calc_draw_probability(draw_margin, size, env)
    if env == nil then
        env = global_env()
    end
    return 2 * env.cdf(draw_margin / (math.sqrt(size) * env.beta)) - 1
end

local function calc_draw_margin(draw_probability, size, env)
    if env == nil then
        env = global_env()
    end
    return env.ppf((draw_probability + 1) / 2) * math.sqrt(size) * env.beta
end

local function v_win(diff, draw_margin)
    local x = diff - draw_margin
    local denom = global_env().cdf(x)
    if denom then 
        return (global_env().pdf(x) / denom) 
    else 
        return -x 
    end
end

local function v_draw(diff, draw_margin)
    local abs_diff = math.abs(diff)
    local a, b = draw_margin - abs_diff, -draw_margin - abs_diff
    local denom = global_env().cdf(a) - global_env().cdf(b)
    local numer = global_env().pdf(b) - global_env().pdf(a)
    if denom then
        return (numer / denom) * (diff < 0 and -1 or 1)
    else
        return a * (diff < 0 and -1 or 1)
    end
end

local function w_win(diff, draw_margin)
    local x = diff - draw_margin
    local v = v_win(diff, draw_margin)
    local w = v * (v + x)
    if 0 < w and w < 1 then
        return w
    end
    error('Cannot calculate correctly, set backend to "mpmath"')
end

local function w_draw(diff, draw_margin)
    local abs_diff = math.abs(diff)
    local a, b = draw_margin - abs_diff, -draw_margin - abs_diff
    local denom = global_env().cdf(a) - global_env().cdf(b)
    if not denom then
        error('Cannot calculate correctly, set backend to "mpmath"')
    end
    local v = v_draw(abs_diff, draw_margin)
    return (v ^ 2) + (a * global_env().pdf(a) - b * global_env().pdf(b)) / denom
end

local function _team_sizes(rating_groups)
    local team_sizes = {}
    local sum = 0
    for _, group in ipairs(rating_groups) do
        sum = sum + #group
        table.insert(team_sizes, sum)
    end
    return team_sizes
end

function TrueSkill:factor_graph_builders(rating_groups, ranks, weights)
    local flatten_ratings = {}
    for _, group in ipairs(rating_groups) do
        for _, rating in ipairs(group) do
            table.insert(flatten_ratings, rating)
        end
    end
    local flatten_weights = {}
    for _, group in ipairs(weights) do
        for _, weight in ipairs(group) do
            table.insert(flatten_weights, weight)
        end
    end
    local size = #flatten_ratings
    local group_size = #rating_groups
    local rating_vars = {}
    for _ = 1, size do table.insert(rating_vars, Variable()) end
    local perf_vars = {}
    for _ = 1, size do table.insert(perf_vars, Variable()) end
    local team_perf_vars = {}
    for _ = 1, group_size do table.insert(team_perf_vars, Variable()) end
    local team_diff_vars = {}
    for _ = 1, group_size - 1 do table.insert(team_diff_vars, Variable()) end
    local team_sizes = _team_sizes(rating_groups)
    local function build_rating_layer()
        for i, rating_var in ipairs(rating_vars) do
            coroutine.yield(PriorFactor(rating_var, flatten_ratings[i], self.tau))
        end
    end
    local function build_perf_layer()
        for i, rating_var in ipairs(rating_vars) do
            coroutine.yield(LikelihoodFactor(rating_var, perf_vars[i], self.beta ^ 2))
        end
    end
    local function build_team_perf_layer()
        for team, team_perf_var in ipairs(team_perf_vars) do
            local start = team == 1 and 1 or team_sizes[team - 1] + 1
            local child_perf_vars = {}
            for i = start, team_sizes[team] do table.insert(child_perf_vars, perf_vars[i]) end
            local coeffs = {}
            for i = start, team_sizes[team] do table.insert(coeffs, flatten_weights[i]) end
            coroutine.yield(SumFactor(team_perf_var, child_perf_vars, coeffs))
        end
    end
    local function build_team_diff_layer()
        for team, team_diff_var in ipairs(team_diff_vars) do
            coroutine.yield(SumFactor(team_diff_var,{team_perf_vars[team], team_perf_vars[team + 1]}, {1, -1}))
        end 
    end 
    local function build_trunc_layer()
        for x, team_diff_var in ipairs(team_diff_vars) do 
            local draw_probability 
            if type(self.draw_probability) == 'function' then 
                draw_probability = self.draw_probability(Rating(team_perf_vars[x]), Rating(team_perf_vars[x + 1]), self)
            else 
                draw_probability = self.draw_probability 
            end 
            local size = #rating_groups[x] + #rating_groups[x + 1]
            local draw_margin = calc_draw_margin(draw_probability, size, self)
            local v_func 
            local w_func 
            if ranks[x] == ranks[x + 1] then 
                v_func = v_draw 
                w_func = w_draw 
            else 
                v_func = v_win 
                w_func = w_win 
            end 
            coroutine.yield(TruncateFactor(team_diff_var,v_func, w_func, draw_margin))
        end 
    end 
    return build_rating_layer, build_perf_layer, build_team_perf_layer,build_team_diff_layer, build_trunc_layer 
end 

function TrueSkill:run_schedule(build_rating_layer, build_perf_layer,build_team_perf_layer, build_team_diff_layer,build_trunc_layer, min_delta)
    min_delta = min_delta or DELTA
    if min_delta <= 0 then
        error('min_delta must be greater than 0')
    end
    local layers = {}
    local function build(builders)
        local layers_built = {}
        for _, build in ipairs(builders) do
            local layer = {}
            for factor in coroutine.wrap(build) do
                table.insert(layer, factor)
            end
            table.insert(layers_built, layer)
            table.insert(layers, layer)
        end
        return layers_built
    end
    local layers_built = build({build_rating_layer,build_perf_layer,build_team_perf_layer})
    local rating_layer, perf_layer, team_perf_layer = table.unpack(layers_built)
    for _, f in ipairs(rating_layer) do f:down() end
    for _, f in ipairs(perf_layer) do f:down() end
    for _, f in ipairs(team_perf_layer) do f:down() end
    local team_diff_layer, trunc_layer = table.unpack(build({build_team_diff_layer,
                                                       build_trunc_layer}))
    local team_diff_len = #team_diff_layer
    for _ = 1, 10 do
        local delta = 0
        if team_diff_len == 1 then
            team_diff_layer[1]:down()
            delta = trunc_layer[1]:up()
        else
            for x = 1, team_diff_len - 1 do
                team_diff_layer[x]:down()
                delta = math.max(delta, trunc_layer[x]:up())
                team_diff_layer[x]:up(1)
            end
            for x = team_diff_len - 1, 1, -1 do
                team_diff_layer[x + 1]:down()
                delta = math.max(delta, trunc_layer[x + 1]:up())
                team_diff_layer[x + 1]:up(0)
            end 
        end 
        if delta <= min_delta then break end 
    end 
    team_diff_layer[1]:up(0)
    team_diff_layer[team_diff_len]:up(1)
    for _, f in ipairs(team_perf_layer) do 
        for x = 1, #f.vars - 1 do 
            f:up(x - 1)
        end 
    end 
    for _, f in ipairs(perf_layer) do f:up() end 
    return layers 
end 

function TrueSkill:rate(rating_groups, ranks, weights, min_delta)
    min_delta = min_delta or DELTA
    rating_groups = self:validate_rating_groups(rating_groups)
    weights = self:validate_weights(weights, rating_groups)
    local group_size = #rating_groups
    if ranks == nil then
        ranks = {}
        for i = 1, group_size do table.insert(ranks, i - 1) end
    elseif #ranks ~= group_size then
        error('Wrong ranks')
    end
    local sorting = {}
    for i, g in ipairs(rating_groups) do
        table.insert(sorting, {i, g, ranks[i], weights[i]})
    end
    table.sort(sorting, function(a, b) return a[3] < b[3] end)
    local sorted_rating_groups, sorted_ranks, sorted_weights = {}, {}, {}
    for _, t in ipairs(sorting) do
        table.insert(sorted_rating_groups, t[2])
        table.insert(sorted_ranks, t[3])
        local w = {}

        for _, w_ in ipairs(t[4]) do table.insert(w, math.max(min_delta, w_)) end
        table.insert(sorted_weights, w)
    end
    local builders = {self:factor_graph_builders(sorted_rating_groups,sorted_ranks,sorted_weights)}
    local build_rating_layer, build_perf_layer, build_team_perf_layer, build_team_diff_layer,build_trunc_layer = table.unpack(builders)
    local layers = {self:run_schedule(build_rating_layer, build_perf_layer, build_team_perf_layer, build_team_diff_layer,build_trunc_layer, min_delta)}
    local rating_layer = layers[1][1]
    local team_sizes = _team_sizes(sorted_rating_groups)
    local transformed_groups = {}
    for i, start in ipairs({0, table.unpack(team_sizes)}) do
        if i > #team_sizes then break end
        local group = {}
        for j = start + 1, team_sizes[i] do
            local f = rating_layer[j]
            table.insert(group, Rating(f:var():mu(), f:var():sigma()))
        end
        table.insert(transformed_groups, group)
    end
    local unsorting = {}
    for i, t in ipairs(sorting) do
        table.insert(unsorting, {t[1], transformed_groups[i]})
    end
    table.sort(unsorting, function(a, b) return a[1] < b[1] end)
    local result = {}
    for _, t in ipairs(unsorting) do table.insert(result, t[2]) end
    return result 
end 

function TrueSkill:quality(rating_groups, weights)
    rating_groups = self:validate_rating_groups(rating_groups)
    local weights = self:validate_weights(weights, rating_groups)
    local flatten_ratings = {}
    for _, group in ipairs(rating_groups) do
        for _, rating in ipairs(group) do
            table.insert(flatten_ratings, rating)
        end
    end
    local flatten_weights = {}
    for _, group in ipairs(weights) do
        for _, weight in ipairs(group) do
            table.insert(flatten_weights, weight)
        end
    end

    local mean = {}
    for i, r in ipairs(flatten_ratings) do
        mean[i] = {r.mu}
    end
    local mean_matrix = Matrix(mean)

    local variance ={}
    for i = 1, #flatten_ratings do
        variance[i] = {}
        for j = 1, #flatten_ratings do
            if i == j then
                variance[i][j] = flatten_ratings[i].sigma ^ 2
            else
                variance[i][j] = 0
            end
        end
    end
    local variance_matrix = Matrix(variance)

    local rotated_a = {}
    local k = 1
    for i = 1, #rating_groups - 1 do
        local cur, _next = rating_groups[i], rating_groups[i + 1]
        rotated_a[i] = rotated_a[i] or {}
        for j = k, k + #cur - 1 do
            rotated_a[i][j] = flatten_weights[j]
        end
        k = k + #cur
        for j = k, k + #_next - 1 do
            rotated_a[i][j] = -flatten_weights[j]
        end
    end
    for i = 1, #rating_groups - 1 do
        for j = 1, #flatten_ratings do
            if not rotated_a[i][j] then
                rotated_a[i][j] = 0
            end
        end
    end
    local rotated_a_matrix = Matrix(rotated_a)

    local a_matrix = rotated_a_matrix:transpose()
    local _ata = rotated_a_matrix * a_matrix * (self.beta ^ 2) 
    local _atsa = rotated_a_matrix * variance_matrix * a_matrix 
    local start = mean_matrix:transpose() * a_matrix 
    local middle = _ata + _atsa 
    local _end = rotated_a_matrix * mean_matrix
    local a = middle:inverse()
    local e_arg = (start * a * _end * -0.5):determinant()
    local s_arg = _ata:determinant() / middle:determinant()
    return math.exp(e_arg) * math.sqrt(s_arg)
end 

function TrueSkill:expose(rating)
    local k = 3
    return rating.mu - k * rating.sigma
end

function TrueSkill:make_as_global()
    setup({env = self})
    return self
end

function TrueSkill:__tostring()
    local c = getmetatable(self)
    local draw_probability
    if type(self.draw_probability) == "function" then
        local f = self.draw_probability
        draw_probability = f.__module .. '.' .. f.__name
    else
        draw_probability = string.format('%.1f%%', self.draw_probability * 100)
    end
    local backend = ''
    if self.backend ~= nil then
        if type(self.backend) == "table" then
            backend = ', backend=...'
        else
            backend = string.format(', backend=%s', tostring(self.backend))
        end
    end
    local args = {c.__module .. '.' .. c.__name, self.mu, self.sigma, self.beta, self.tau, draw_probability, backend}
    return string.format('%s(mu=%.3f, sigma=%.3f, beta=%.3f, tau=%.3f, draw_probability=%s%s)', table.unpack(args))
end

-- 设置元表，使得可以使用TrueSkill()语法创建新对象
setmetatable(TrueSkill, {__call = TrueSkill.new})

local function rate(rating_groups, ranks, weights, min_delta)
    return global_env():rate(rating_groups, ranks, weights, min_delta)
end

local function quality(rating_groups, weights)
    return global_env():quality(rating_groups, weights)
end

local function expose(rating)
    return global_env():expose(rating)
end

return {
    Rating = Rating,
    rate = rate,
    quality = quality,
    expose = expose,
}